package com.xiao.xin.core.service.Impl;

import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson2.JSON;
import com.xiao.xin.core.Util.UserThreadLocal;
import com.xiao.xin.core.common.Constants;
import com.xiao.xin.core.domain.*;
import com.xiao.xin.core.dto.CancelComfyuiTaskDto;
import com.xiao.xin.core.dto.Text2ImageDto;
import com.xiao.xin.core.exception.Text2ImageExecption;
import com.xiao.xin.core.exception.UserFundException;
import com.xiao.xin.core.service.*;
import com.xiao.xin.core.vo.Text2ImageVo;
import com.xiao.xin.core.vo.UserResultVo;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Service;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
@Service
@RequiredArgsConstructor
public class Text2ImageServiceImpl implements Text2ImageService {

    private final OllamaChatService ollamaChatService;
    private final FreemarkerService freemarkerService;
    private final RedisService redisService;
    private final UserFundRecordService userFundRecordService;
    private final UserResultService userResultService;
    private final StringRedisTemplate redisTemplate;

    @Override
    public ComfyuiTask getComfyuiTask(Text2ImageDto text2ImageDto) throws Exception {

        String propmt = null;
        if (StrUtil.isNotBlank(text2ImageDto.getPropmt())) {
            propmt = ollamaChatService.translate(text2ImageDto.getPropmt());
        }

        String reverse = null;
        // AI 翻译
        if (StrUtil.isNotBlank(text2ImageDto.getReverse())) {
            reverse = ollamaChatService.translate(text2ImageDto.getReverse());
        }

        // 配置模板数据
        ComfyuiModel comfyuiModel = ComfyuiModel
                .builder()
                .seed(text2ImageDto.getSeed())
                .step(text2ImageDto.getStep())
                .cfg(text2ImageDto.getCfg())
                .size(text2ImageDto.getSize())
                .modelName(text2ImageDto.modelName())
                .width(text2ImageDto.width())
                .height(text2ImageDto.height())
                .propmt(("(8k, best quality, masterpiece),(high detailed skin),") + propmt)
                .samplerName(text2ImageDto.samplerName())
                .scheduler(text2ImageDto.scheduler())
                .reverse(reverse + ",bad face,naked,bad finger,bad arm,bad leg,bad eye")
                .build();

        log.info("\n正向词:{}\n反向词:{}", comfyuiModel.getPropmt(), comfyuiModel.getReverse());

        // 修改模板数据
        String promp = freemarkerService.renderText2Image(comfyuiModel);

        log.info("模板数据:{}", promp);

        // 配置ComfyUI添加队列数据
        ComfyuiRequestDto comfyuiRequestDto = ComfyuiRequestDto
                .builder()
                .clientId(Constants.COMFYUI_CLIENT_ID)
                .prompt(JSON.parseObject(promp))
                .build();

        // 配置ComfyUITask 数据
        ComfyuiTask comfyuiTask = new ComfyuiTask(text2ImageDto.getClientId(), comfyuiRequestDto);
        comfyuiTask.setUserId(UserThreadLocal.getUser().getId());
        comfyuiTask.setSize(text2ImageDto.getSize());
        return comfyuiTask;
    }


    @Override
    public Text2ImageVo propmt(Text2ImageDto dto) throws Exception {
        if (dto.getSize() < 1) {
            throw new UserFundException("请求参数错误");
        }

        // 冻结积分
        userFundRecordService.pointsFreeze(UserThreadLocal.getUser().getId(), dto.getSize());

        // 获取ComfyuiTask封装好的对象
        ComfyuiTask comfyuiTask = getComfyuiTask(dto);


        // 添加Redis到ZSet有序集合中
        redisService.addQueueTask(comfyuiTask);


        return Text2ImageVo
                .builder()
                .pid(comfyuiTask.getPromptId())
                .queueIndex(comfyuiTask.getIndex())
                .build();
    }

    @Override
    public void cancel(CancelComfyuiTaskDto cancelComfyuiTaskDto) {
        String IsStart = redisTemplate.opsForValue().get(Constants.START_TASK);
        if (StrUtil.isNotEmpty(IsStart)) {
            throw new Text2ImageExecption("任务已经开始或不存在");
        }

        // 获取当前任务 ComfyUiTask
        ComfyuiTask comfyuiTask = redisService.getStartedTask(cancelComfyuiTaskDto.getTempId());

        log.info("取消任务获取当前任务:{}",comfyuiTask);

        // 执行删除
        redisTemplate.opsForZSet().remove(Constants.DISTRIBUTED_QUEUE_KEY, comfyuiTask.getId());
        redisTemplate.delete(Constants.TASK_KEY_PREFIX + comfyuiTask.getId());

        // 归还用户积分
        userFundRecordService.freezeReturn(comfyuiTask.getUserId(), comfyuiTask.getSize());
    }

    @Override
    public Map list(QueryTasklist queryTasklist) {
        Map map = userResultService.listUserResult(queryTasklist);
        return map;
    }
}
